#include "StdAfx.h"
#include "MaryClient.h"

#pragma warning(disable : 4996)

/**
* A C++ implementation of a simple client to the MARY TTS system.
* server_host: string representation of the host where the server is running
* server_port: the TCP port where the server is running
*
*/
MaryClient::MaryClient(const string& _server_host, int _server_port)
{
	// Initialise network stack interface
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2,2), &wsaData);

	server_host = _server_host;
	server_port = _server_port;
	connected = false;
	infoSent = false;


	// set default types
	setTypes("TEXT","AUDIO","WAVE","female","");

}

MaryClient::~MaryClient(void)
{
	// clean up network stack interface
	WSACleanup();

}


/**
* Setter methods
*/
int MaryClient::setTypes(const string& _inputType, const string& _outputType, const string& _audioType, const string& _voiceName, const string& _effects)
{

	inputType = _inputType;
	outputType = _outputType;
	audioType = _audioType;
	voiceName = _voiceName;
	effects = _effects;
	locale = "en-GB";

	return ERROR_NONE;

}

int MaryClient::setInputType(const string& _inputType)
{
	inputType = _inputType;
	return ERROR_NONE;
}


int MaryClient::setOutputType(const string& _outputType)
{
	outputType = _outputType;
	return ERROR_NONE;
}

int MaryClient::setAudioType(const string& _audioType)
{
	audioType = _audioType;
	return ERROR_NONE;
}

int MaryClient::setVoiceName(const string& _voiceName)
{	
	voiceName = _voiceName;
	if (voiceName == "")
	{
		voiceName = voiceList[0]->GetName();
	}

	// set locale of chosen voice
	locale = "";
	for (size_t v = 0; v < voiceList.size(); v++)
	{
		if (voiceName == voiceList[v]->GetName())
		{
			locale = voiceList[v]->GetLocale();
			break;
		}
	}


	return ERROR_NONE;
}

int MaryClient::setEffects(const string& _effects)
{
	effects = _effects;
	return ERROR_NONE;
}

/**
* Getter methods
*/
string MaryClient::getOutputType()
{
	return outputType;
}

int MaryClient::process(const string& inputText, string& result)
{

	return process(inputText, this->inputType, this->outputType, this->locale, this->audioType, this->voiceName, this->effects, result);

}

int MaryClient::connectInfoSocket()
{	
	// create a tcp connection to the mary server
	maryInfoSocket = socket(AF_INET, SOCK_STREAM, 0);

	// verify that the socket could be opened successfully
	if (maryInfoSocket == -1)
	{
		return ERROR_INFO_SOCKET_CREATE;
	}
	else // autoflush stdout, bind and connect		
	{
		maryClient.sin_family = AF_INET;
		maryClient.sin_port = htons (0);
		maryClient.sin_addr.s_addr = INADDR_ANY;

		int status = bind (maryInfoSocket, (struct sockaddr*) &maryClient, sizeof (maryClient));

		if (status != 0)
		{
			return ERROR_INFO_SOCKET_BIND;
		}

		maryServer.sin_family = AF_INET;
		maryServer.sin_port = htons(static_cast<u_short>(server_port));
		memcpy ((char*) &maryServer.sin_addr.s_addr, hostInfo->h_addr_list [0], hostInfo->h_length);

		status = connect (maryInfoSocket, (struct sockaddr*) &maryServer, sizeof (maryServer));

		if (status != 0)
		{
			return ERROR_INFO_SOCKET_CONNECT;
		}
	}

	return ERROR_NONE;
}

int MaryClient::connectDataSocket()
{
	// create a tcp connection to the mary server
	maryDataSocket = socket(AF_INET, SOCK_STREAM, 0);

	// verify that the socket could be opened successfully
	if (maryDataSocket == -1)
	{
		return ERROR_DATA_SOCKET_CREATE;
	}
	else // autoflush stdout, bind and connect		
	{
		maryClient.sin_family = AF_INET;
		maryClient.sin_port = htons (0);
		maryClient.sin_addr.s_addr = INADDR_ANY;

		int status = bind(maryDataSocket, (struct sockaddr*) &maryClient, sizeof (maryClient));

		if (status != 0)
		{
			return ERROR_DATA_SOCKET_BIND;
		}

		maryServer.sin_family = AF_INET;
		maryServer.sin_port = htons(static_cast<u_short>(server_port));
		memcpy ((char*) &maryServer.sin_addr.s_addr, hostInfo->h_addr_list [0], hostInfo->h_length);

		status = connect(maryDataSocket, (struct sockaddr*) &maryServer, sizeof (maryServer));

		if (status != 0)
		{
			return ERROR_DATA_SOCKET_CONNECT;
		}
	}
	return ERROR_NONE;
}

int MaryClient::sendInfoRequest(const string& maryInFormat,
								const string& maryOutFormat,
								const string& locale,
								const string& audioType,
								const string& voiceName,
								const string& effects,
								char* id)
{

    infoSent = false;

	// prepare the request
	string query = "MARY IN=" + maryInFormat + " OUT=" + maryOutFormat + " LOCALE=" + locale;

	if (maryOutFormat == "AUDIO") query += " AUDIO="+audioType;
	
	if (voiceName != "")  
	{
		//query += " LOCALE=" + locale;	// remove this line, if using an older version than MARY 4.0
		query += " VOICE="+voiceName;
	}

	if (effects != "") query += " EFFECTS="+ effects;

	query += "\r\n";

	// send request to the Mary server
	if (send(maryInfoSocket, query.c_str (), query.size (), 0) == -1)
	{
		return ERROR_INFO_SENDING_REQUEST;
	}

	// receive the request id
	if (recv(maryInfoSocket, id, 32, 0) == -1)
	{
		return ERROR_INFO_RECEIVING_ID;
	}

	// check an integer is returned from the server
	if(!atoi(id))
	{
		return ERROR_INFO_MALFORMED_REQUEST;
	}


	infoSent = true;
	return ERROR_NONE;

}


int MaryClient::sendDataRequest(const string& inputText, char* id)
{

	// send the request id to the Mary server
	if (send(maryDataSocket, id, strlen (id), 0) == -1)
	{
		return ERROR_DATA_SENDING_ID;
	}

	// send the query to the Mary server
	if (send(maryDataSocket, inputText.c_str(), inputText.size(), 0) == -1)
	{
		return ERROR_DATA_SENDING_TEXT;
	}

	// send terminating charaters
	if (send(maryDataSocket, "\r\n", 2, 0) == -1)
	{
		return ERROR_DATA_SENDING_TEXT;
	}


	return ERROR_NONE;
}

int MaryClient::connectToServer()
{
	int inError;

	// get host information
	hostInfo = gethostbyname (server_host.c_str());

	if (hostInfo == NULL)
	{
		return -2;
	}

	inError = connectInfoSocket();
	if (inError) return inError;

	inError = connectDataSocket();
	if (inError) return inError;

	this->connected = true;

	return ERROR_NONE;
}


int MaryClient::getVoices(vector<string>& voices)
{

	int inError = 0;
	const int buflen = 1024;
	char recvbuf[buflen] = "";

	// clear voice vector
	for (size_t v = 0; v < voiceList.size(); v++)
	{
		delete voiceList[v];
	}

	inError = connectToServer();
	if (inError) return inError;


	// prepare the request
	string query = "MARY LIST VOICES\r\n";

	// send request to the Mary server
	if (send(maryInfoSocket, query.c_str (), query.size (), 0) == -1)
	{
		return ERROR_INFO_SENDING_REQUEST;
	}

	// receive the request information
	int iResult;
	bool done = false;
	do
	{
		iResult = recv(maryInfoSocket, recvbuf, buflen, 0);		

		if (iResult == SOCKET_ERROR)
		{
			return ERROR_INFO_RECEIVING_ID;
		}
		
		// parse the char array returned from Mary (list of names, separated by newline, ended by empty line
		char *token = NULL;
		token = strtok(recvbuf,"\n");

		while (token != NULL)
		{	
			// check for empty line
			if (token[0] == 13)
			{
				done = true;
				break;
			}

			// get the first symbol (the voice name)
			MaryVoice* voice = new MaryVoice(string(token));
			if (voice->IsValid())
			{
				voiceList.push_back(voice);
				voices.push_back(voice->GetName());
			}

			// get the next token
			token = strtok( NULL, "\n" );
			

		}

	} while( iResult > 0  && !done);

	return ERROR_NONE;

}

/**
* This code is modified from various places.


* inputText: the UTF-8 encoded text (or XML document) to send as a request
* maryInFormat: the input type of the data in inputText, e.g. TEXT_DE
* maryOutFormat: the output type to produce, e.g. MBROLA, AUDIO
* audioType: for AUDIO output, the type of audio data to produce,
*            e.g. WAVE or MP3
* voiceName: the name of the voice to use, e.g. female
* result: an empty string serving as the container for the output.
*
*    It will return text or audio data; text data will be encoded as UTF-8.
*
* return value: 0 on success, negative on failure.
*
*
*/
int MaryClient::process(const string& inputText,
						const string& maryInFormat,
						const string& maryOutFormat,
						const string& locale,
						const string& audioType,
						const string& voiceName,
						const string& effects,
						string& result)
{

	int inError = 0;

	inError = connectToServer();
	if (inError) return inError;

	char id[32] = "";
	
	inError = sendInfoRequest(maryInFormat,maryOutFormat,locale,audioType,voiceName,effects,id);
	if (inError) return inError;
	

	inError = sendDataRequest(inputText, id);
	if (inError) return inError;

	// shutdown data socket
	shutdown(maryDataSocket, SD_SEND);

	// get the result from the mary server
	inError = readandfill(maryDataSocket, result);
	if (inError) return inError;

	// close all open sockets
	closesocket (maryInfoSocket);
	closesocket (maryDataSocket);
	this->connected = false;

	return ERROR_NONE;
}




int MaryClient::readandfill(int maryDataSocket, string& result)
{
	unsigned int total_bytes = 0;
	int recv_bytes = 0;
	char data [2048] = "";

	result.clear();
	result[0] = '\0';

	// receive the request result
	do
	{
		data [0] = '\0';

		recv_bytes = recv(maryDataSocket, data, 1024, 0);

		if (recv_bytes == -1)
		{
			return ERROR_DATA_RECEIVING_DATA;
		}
		else if (recv_bytes > 0)
		{
			//cerr << "("<<recv_bytes<<")";
			total_bytes += recv_bytes;
			data [recv_bytes] = '\0';

			if (this->outputType == "AUDIO")
			{
				for (int i=0; i<recv_bytes; i++)
				{
					result += data [i];
				}
			}
			else
			{
				result += data;
			}
		}
	} while (recv_bytes != 0);

	if (result.size () != total_bytes)
	{
		cerr << "error: total bytes received != result bytes!" << endl;
		cerr << "       total bytes received = " << total_bytes << endl;
		cerr << "       result bytes = " << result.size () << endl;
	}

	return ERROR_NONE;
}

/**
 * Returns a string description of the error
 */
string MaryClient::getErrorString(int errorCode)
{

	ostringstream output;

	switch (errorCode)
	{
	case ERROR_NONE:
		output << "No Error.";
		break;
	case ERROR_INFO_SOCKET_CONNECT:
	case ERROR_DATA_SOCKET_CONNECT:
		output << "Cannot connect to Mary server: " << this->server_host << ":" << this->server_port << "!";
		break;
	case ERROR_INFO_SOCKET_CREATE:
	case ERROR_DATA_SOCKET_CREATE:
	case ERROR_INFO_SOCKET_BIND:
	case ERROR_DATA_SOCKET_BIND:
		output << "Cannot create or bind socket!";
		break;
	case ERROR_INFO_SENDING_REQUEST:
		output << "Could not send info request.";
		break;
	case ERROR_INFO_RECEIVING_ID:
		output << "Could not receive ID.";
		break;
	case ERROR_INFO_MALFORMED_REQUEST:
		output << "Malformed Information Request!";
		break;
	case ERROR_DATA_SENDING_ID:
		output << "Could not send ID for data request.";
		break;
	case ERROR_DATA_SENDING_TEXT:
		output << "Could not send data request.";
		break;
	case ERROR_DATA_RECEIVING_DATA:
		output << "No data to receive!";
		break;
	case ERROR_PLAY_READ_WAVE_FAIL:
		output << "Could not read wave header and data!";
		break;
	case ERROR_PLAY_OPEN_DEVICE_FAIL:
		output << "Could not open device to play wave data!";
		break;
	case ERROR_PLAY_WRITE_DEVICE_FAIL:
		output << "Could not write to audio device!";
		break;
	default :
		output <<  "Unknown error code: " << errorCode << "!";
	}

	return output.str();
}



/**
* Play the result, assuming that result contains audio!
*/
int MaryClient::play(string& result, HWAVEOUT &hWaveOut, WAVEHDR &waveHdr, double* factor)
{

	string error;
	UINT  wResult;
	WAVEFORMATEX *  pWaveHeader = 0;
	void *  pbWaveData = 0;
	DWORD  cbWaveSize = 0;

	// get the wave data from the input result
	BOOL ok = GetWaveData((void *)result.c_str(), pWaveHeader, pbWaveData, cbWaveSize);

	if (!ok) return ERROR_PLAY_READ_WAVE_FAIL;
	

	*factor = ((double)pWaveHeader->nSamplesPerSec * pWaveHeader->wBitsPerSample / 8);
	

    // Open a waveform device for output using window callback. 
    if (waveOutOpen((LPHWAVEOUT)&hWaveOut, WAVE_MAPPER, pWaveHeader, 0L, 0L, CALLBACK_NULL)) 
    { 
		//error = "Failed to open waveform output device.";
        return ERROR_PLAY_OPEN_DEVICE_FAIL; 
    } 

	// Alloc the sample buffer
    memset(&waveHdr,0,sizeof(WAVEHDR));   
   
	//dwDataSize = result.length();
 
    // After allocation, set up and prepare header. 
	waveHdr.lpData = (LPSTR)pbWaveData;
    waveHdr.dwBufferLength = cbWaveSize; 
    waveHdr.dwFlags = 0L; 
    waveHdr.dwLoops = 0L; 
    waveOutPrepareHeader(hWaveOut, &waveHdr, sizeof(WAVEHDR)); 
 
    // Now the data block can be sent to the output device. The 
    // waveOutWrite function returns immediately and waveform 
    // data is sent to the output device in the background. 
    wResult = waveOutWrite(hWaveOut, &waveHdr, sizeof(WAVEHDR)); 
    if (wResult != 0) 
    { 
        waveOutUnprepareHeader(hWaveOut, &waveHdr, sizeof(WAVEHDR)); 
        error =  "Failed to write block to device"; 
		return ERROR_PLAY_WRITE_DEVICE_FAIL; 
    } 
	return 0;
	

}

//----------------------------------------------------------------------------------------------------

BOOL MaryClient::GetWaveData(void * pRes, WAVEFORMATEX * & pWaveHeader, void * & pbWaveData, DWORD & cbWaveSize) 
{
	pWaveHeader = 0;
	pbWaveData = 0;
	cbWaveSize = 0;

	DWORD * pdw = (DWORD *)pRes;
	DWORD dwRiff = *pdw++;
	DWORD dwLength = *pdw++;
	DWORD dwType = *pdw++;

	if( dwRiff != mmioFOURCC('R', 'I', 'F', 'F') )
		return FALSE;      // not even RIFF

	if( dwType != mmioFOURCC('W', 'A', 'V', 'E') )
		return FALSE;      // not a WAV

	DWORD * pdwEnd = (DWORD *)((BYTE *)pdw + dwLength-4);

	while( pdw < pdwEnd ) {
		dwType = *pdw++;
		dwLength = *pdw++;

		switch( dwType ) {
			case mmioFOURCC('f', 'm', 't', ' '):
				if( !pWaveHeader ) {
					if( dwLength < sizeof(WAVEFORMAT) )
						return FALSE;      // not a WAV

					pWaveHeader = (WAVEFORMATEX *)pdw;

					if( pbWaveData && cbWaveSize )
						return TRUE;
				}
				break;

			case mmioFOURCC('d', 'a', 't', 'a'):
				pbWaveData = LPVOID(pdw);
				cbWaveSize = dwLength;

				if( pWaveHeader )
					return TRUE;
				break;
		}
		pdw = (DWORD *)((BYTE *)pdw + ((dwLength+1)&~1));
	}

	return FALSE;
}

/*******************************************************
************* MaryVoice Implementation *****************
********************************************************/

MaryVoice::MaryVoice(const string &voiceString)
{
	this->source = voiceString;

	// parse something like dfki-spike en_GB male unitselection general	
	size_t namei,locali;

	namei = source.find(" ");
	if (namei!=string::npos)
	{
		this->name = source.substr(0,namei);
		locali = source.find(" ",namei+1,1);
		if (locali!=string::npos)
		{
			this->locale = source.substr(namei+1,locali-namei);
		}
	}

	this->valid = true;
}

MaryVoice::~MaryVoice()
{

}

bool MaryVoice::IsValid()
{
	return this->valid;
}

string MaryVoice::GetName()
{
	return this->name;
}

string MaryVoice::GetLocale()
{
	return this->locale;
}


